#%%
import pickle
import numpy as pd
import os
import os.path as osp
import sys
import shutil
import h5py
import numpy as np
import pickle
import cv2
import pandas as pd
import tqdm
import multiprocessing as mp
from lilab.openlabcluster_postprocess.s1a_clipNames_inplace_parse import parse_name

iter_n = 1
segLength = 24
project_iter0 = '/DATA/taoxianming/rat/data/Mix_analysis/SexAgeDay55andzzcWTinAUT_MMFF/result32/feats32-2024-04-11'
video_src_400p = '/DATA/taoxianming/rat/data/Mix_analysis/SexAgeDay55andzzcWTinAUT_MMFF/data'
project_out = osp.dirname(project_iter0) + '/result32_iter{}'.format(iter_n)
predpklfile = '/home/liying_lab/chenxf/ml-project/论文图表/semisupervised分类/MM_FF_Representive_K34.clippredpkl'
predpkldata = pickle.load(open(predpklfile, 'rb'))
ind_rawclip = predpkldata['ind_rawclip']
cluster_labels_sub = predpkldata['cluster_labels']
clipNames_sub = predpkldata['clipNames']
cluster_labels_sub = cluster_labels_sub - cluster_labels_sub.min() + 1 #start from 1

video_dst_400p = osp.dirname(project_iter0) + '/video_file'

def create_copy_project():
    os.makedirs(video_dst_400p, exist_ok=True)
    os.makedirs(project_out, exist_ok=True)
    subdirs = ['datasets', 'label', 'models', 'models_Dec', 'output', 'sample', 'videos']
    for subdir in subdirs:
        os.makedirs(osp.join(project_out, subdir), exist_ok=True)

    shutil.copyfile(osp.join(project_iter0, 'videos/clipNames.txt'), osp.join(project_out, 'videos/clipNames.txt'))
    shutil.copyfile(osp.join(project_iter0, 'datasets/data.h5'), osp.join(project_out, 'datasets/data.h5'))
    shutil.copyfile(osp.join(project_iter0, 'config.yaml'), osp.join(project_out, 'config.yaml'))

clipNames_file_out = osp.join(project_out, 'videos/clipNames.txt')
clipNames = np.array([osp.basename(f.strip()) for f in (open(clipNames_file_out, 'r').readlines())])

create_copy_project()

def update_data_h5():
    data_h5_out = osp.join(project_out, 'datasets/data.h5')
    assert clipNames.dtype == clipNames_sub.dtype and np.all(clipNames[ind_rawclip] == clipNames_sub)

    with h5py.File(data_h5_out,'r+') as hf:
        data = np.zeros(len(clipNames), dtype=int)
        data[ind_rawclip] = cluster_labels_sub
        a = hf.pop('label')
        hf.create_dataset('label',data=data)

    np.save(osp.join(project_out, 'label/label.npy'), data)

update_data_h5()


df_clipNames = parse_name(clipNames)
df_clipNames['clipName'] = clipNames
df_clipNames_B = df_clipNames.loc[df_clipNames['isBlack']][['vnake', 'startFrame', 'clipName']].sort_values(['vnake', 'startFrame'])
df_clipNames_W = df_clipNames.loc[~df_clipNames['isBlack']][['vnake', 'startFrame', 'clipName']].sort_values(['vnake', 'startFrame'])
assert np.all(df_clipNames_B[['vnake', 'startFrame']].values == df_clipNames_W[['vnake', 'startFrame']].values)
vnakes = df_clipNames['vnake'].unique()

assert all(osp.isfile(osp.join(video_src_400p, vnake_now+'_1_400p.mp4')) for vnake_now in vnakes)

def create_video_clips(vnake_now):
    src_400p = osp.join(video_src_400p, vnake_now+'_1_400p.mp4')
    assert osp.isfile(src_400p)
    cap = cv2.VideoCapture(src_400p)
    # get fps,frameSize of cap
    frameSize = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
    fps = cap.get(cv2.CAP_PROP_FPS)
    startFramesB, clipNamesB = df_clipNames_B.loc[df_clipNames_B['vnake'] == vnake_now][['startFrame','clipName']].values.T
    startFramesW, clipNamesW = df_clipNames_W.loc[df_clipNames_W['vnake'] == vnake_now][['startFrame','clipName']].values.T
    assert np.all(startFramesB == startFramesW)
    videos_dst_400p_B = [osp.join(video_dst_400p, v) for v in clipNamesB]
    videos_dst_400p_W = [osp.join(video_dst_400p, v) for v in clipNamesW]
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')

    for sf, video_dst_400p_B, video_dst_400p_W in zip(tqdm.tqdm(startFramesB), videos_dst_400p_B, videos_dst_400p_W):
        ##set start
        cap.set(cv2.CAP_PROP_POS_FRAMES, sf)
        outB = cv2.VideoWriter(video_dst_400p_B, fourcc, fps,frameSize)
        outW = cv2.VideoWriter(video_dst_400p_W, fourcc, fps,frameSize)
        for _ in range(segLength):
            success, frame = cap.read()
            frameR = frame.copy()
            outB.write(frameR)
            frameR = cv2.putText(frameR, 'W', (50, 100), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 125, 255), 2)
            outW.write(frameR)
        outB.release()
        outW.release()
    cap.release() 


#%%
# for vnake_now in vnakes: create_video_clips(vnake_now)
# change this to use multiprocessing
pool = mp.Pool(processes=12)
pool.map(create_video_clips, vnakes)